# -*- coding: utf-8 -*-
"""
参数格式转换工具 - 将扁平格式转换为嵌套格式

@author: hsy
"""
from typing import Dict, Any
from loguru import logger


def convert_to_update_format(parameter_changes: Dict[str, Any], session_id: str) -> Dict[str, Any]:
    """
    将扁平格式转换为更新接口需要的嵌套格式
    
    Args:
        parameter_changes: 扁平格式的参数修改
            例如: {"initial_state.water_level": 100.0, "parameters.B": 0.4}
        session_id: 会话ID
        
    Returns:
        嵌套格式的参数修改
        例如: {
            "session_id": "xxx",
            "updates": {
                "initialState": {"water_level": 100.0},
                "parameters": {"B": 0.4}
            }
        }
    """
    try:
        initial_state = {}
        parameters = {}
        
        # 遍历参数修改，按类别分类
        for key, value in parameter_changes.items():
            if key.startswith("initial_state."):
                # 提取参数名（去掉前缀）
                param_name = key.replace("initial_state.", "")
                initial_state[param_name] = value
            elif key.startswith("parameters."):
                # 提取参数名（去掉前缀）
                param_name = key.replace("parameters.", "")
                parameters[param_name] = value
            else:
                # 未知格式的参数，记录警告
                logger.warning(f"未知参数格式: {key}, 将被忽略")
        
        # 构建 updates 字典
        updates = {}
        if initial_state:
            updates["initialState"] = initial_state
        if parameters:
            updates["parameters"] = parameters
        
        # 构建最终请求格式
        result = {
            "session_id": session_id,
            "updates": updates
        }
        
        logger.debug(f"参数格式转换完成: {len(initial_state)} 个初始状态, {len(parameters)} 个模型参数")
        return result
        
    except Exception as e:
        logger.error(f"参数格式转换失败: {str(e)}", exc_info=True)
        # 返回空的更新格式
        return {
            "session_id": session_id,
            "updates": {}
        }


# 使用示例
if __name__ == "__main__":
    # 测试转换函数
    test_changes = {
        "initial_state.water_level": 100.10,
        "initial_state.UAO": 1.0,
        "parameters.B": 0.4,
        "parameters.WUM": 0.21
    }
    
    result = convert_to_update_format(test_changes, "test_session_123")
    
    print("原始格式:")
    print(test_changes)
    print("\n转换后格式:")
    import json
    print(json.dumps(result, indent=2, ensure_ascii=False))

